import torch
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader


class Module(nn.Module):
    def __init__(self):
        super(Module,self).__init__()
        self.models = nn.Sequential(nn.Conv2d(in_channels=3,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=2),
                                    nn.Conv2d(in_channels=32,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Conv2d(in_channels=32,out_channels=64,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Flatten(),
                                    nn.Linear(1024,64),
                                    nn.Linear(64,10))

    def forward(self,x):
        x = self.models(x)
        return x


data_test = torchvision.datasets.CIFAR10("./datavision", train=False, transform=torchvision.transforms.ToTensor())
datas = DataLoader(data_test,batch_size=64,shuffle=True,drop_last=True)
loss = nn.CrossEntropyLoss()
MyModule = Module()
optim = torch.optim.SGD(params=MyModule.parameters(),lr=0.01)
for epoch in range(20):
    running_result = 0.0
    for data in datas:
        imgs,targets = data
        output = MyModule(imgs)
        result_loss = loss(output,targets)
        optim.zero_grad()
        result_loss.backward()
        optim.step()
        running_result += result_loss
    print(running_result)
